﻿//-----------------------------------------------------------------------------
// <copyright company="Cortego AB, 2010" file="FIMAttributeStore.cs">
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND.
// </copyright>
//
// Visit this link for more information about custom attribute stores.
// http://msdn.microsoft.com/en-us/library/ee895358.aspx
//-----------------------------------------------------------------------------

namespace Cortego.ADFS.FIMAttributeStore
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Linq;
    using System.Net;
    using System.Security.Principal;
    using System.ServiceModel;
    using System.Xml.Schema;
    using Microsoft.IdentityModel.Threading;
    using Microsoft.IdentityServer.ClaimsPolicy.Engine.AttributeStore;
    using Microsoft.ResourceManagement.Client;
    using Microsoft.ResourceManagement.Client.WsEnumeration;
    using Microsoft.ResourceManagement.Client.WsTransfer;
    using Microsoft.ResourceManagement.ObjectModel;
    using RemoteMessagingAsyncResult = System.Runtime.Remoting.Messaging.AsyncResult;
    
    /// <summary>
    /// Custom ADFS Attribute Store that enables ADFS to fetch attributes from Forefront Identity Manager.
    /// </summary>
    public class FIMAttributeStore : IAttributeStore, IDisposable
    {
        #region Constants and Variables

        // Config parameter name constants etc.
        private const string Endpoint = "Endpoint";
        private const string FIMServiceSPN = "FIMServiceSPN";
        private const string TracingEnabled = "TracingEnabled";
        private const string UserName = "UserName";
        private const string UserDomain = "UserDomain";
        private const string Password = "Password";
        private const string LogSource = "Cortego.ADFS.FIMAttributeStore";
        
        // Config parameter value variables etc.
        private bool tracingEnabled = false;
        private bool disposed = false;
        
        // Object variables.
        private WsEnumerationClient enumerationClient;
        private MexClient mexClient;
        private RmRequestFactory requestFactory;
        private RmResourceFactory resourceFactory;
        private EventLogTraceListener traceListener;
        
        #endregion

        /// <summary>
        /// Initializes a new instance of the FIMAttributeStore class.
        /// </summary>
        public FIMAttributeStore() 
        { 
        }

        /// <summary>
        /// Finalizes an instance of the FIMAttributeStore class.
        /// </summary>
        ~FIMAttributeStore()
        {
            this.Dispose(false);
        }

        /// <summary>
        /// Delegate for calling for query execution method.
        /// </summary>
        /// <param name="formattedQuery">Query coming from ADFS.</param>
        /// <returns>The results of the asynchronous query against FIM.</returns>
        private delegate string[][] RunQueryDelegate(string formattedQuery);

        #region IAttributeStore members

        /// <summary>
        /// Initializes Attribute Store for coomunicating with Forefront Identity Manager. 
        /// </summary>
        /// <param name="config">Configuration settings from ADFS.</param>
        /// <exception cref="Exception">Throws generic Exception on any exception that occurs during Initialize.</exception>
        public void Initialize(Dictionary<string, string> config)
        {
            // Get Configuration settings.
            if (config.ContainsKey(TracingEnabled))
            {
                bool.TryParse(config[TracingEnabled], out this.tracingEnabled);
            }

            string fimServiceSPN = config.ContainsKey(FIMServiceSPN) ? config[FIMServiceSPN] : null;
            string userName = config.ContainsKey(UserName) ? config[UserName] : null;
            string userDomain = config.ContainsKey(UserDomain) ? config[UserDomain] : null;
            string password = config.ContainsKey(Password) ? config[Password] : null;
            string endpoint = config.ContainsKey(Endpoint) ? config[Endpoint] : "http://localhost:5725";

            if (endpoint.EndsWith("/"))
            {
                endpoint = endpoint.Remove(endpoint.Length - 1);
            }

            try
            {
                if (this.tracingEnabled)
                {
                    // Initialize TraceListener.
                    if (Trace.Listeners[LogSource] == null)
                    {
                        this.traceListener = new EventLogTraceListener(LogSource);
                        this.traceListener.TraceOutputOptions |= TraceOptions.DateTime | TraceOptions.ProcessId | TraceOptions.ThreadId | TraceOptions.Callstack;
                        this.traceListener.Name = LogSource;
                        Trace.Listeners.Add(this.traceListener);
                    }

                    string str = string.Format("\nEndpoint: {0}\nFimServiceSPN: {1}\nUserName: {2}\nUserDomain: {3}\nVerboseLogging: {4}", endpoint, fimServiceSPN, userName, userDomain, this.tracingEnabled);
                    Trace.Write("Using the following configuration: " + str);
                }

                // Initilize MEX binding.
                WSHttpBinding mexBinding = new WSHttpBinding(SecurityMode.None);
                mexBinding.Security.Transport.ClientCredentialType = HttpClientCredentialType.Windows;
                mexBinding.MaxReceivedMessageSize = 1048576;

                // Initialize WSEnumeration binding
                WSHttpContextBinding enumerationBinding = new WSHttpContextBinding();
                enumerationBinding.Security.Message.EstablishSecurityContext = false;
                enumerationBinding.Security.Transport.ClientCredentialType = HttpClientCredentialType.Windows;
                enumerationBinding.ContextProtectionLevel = System.Net.Security.ProtectionLevel.Sign;
                enumerationBinding.MaxReceivedMessageSize = 1048576;

                // Initialize endpoints (MEX and WSEnum).
                Uri mexEndpointUri = new Uri(string.Format("{0}/ResourceManagementService/MEX", endpoint));
                Uri enumerationEndpointUri = new Uri(string.Format("{0}/ResourceManagementService/Enumeration", endpoint));
                EndpointAddress mexEndpoint = null;
                EndpointAddress enumerationEndpoint = null;
                if (string.IsNullOrEmpty(fimServiceSPN))
                {
                    mexEndpoint = new EndpointAddress(mexEndpointUri);
                    enumerationEndpoint = new EndpointAddress(enumerationEndpointUri);

                    if (this.tracingEnabled)
                    {
                        Trace.Write(String.Format("Configured endpoints without SPN Identity claim.\nFIM MEX Endpoint: {0}\nFIM Enumeration Endpoint: {1}", mexEndpointUri.AbsoluteUri, enumerationEndpointUri.AbsoluteUri));
                    }
                }
                else
                {
                    SpnEndpointIdentity spnIdentity = new SpnEndpointIdentity(fimServiceSPN);
                    mexEndpoint = new EndpointAddress(mexEndpointUri, spnIdentity, new System.ServiceModel.Channels.AddressHeader[] { });
                    enumerationEndpoint = new EndpointAddress(enumerationEndpointUri, spnIdentity, new System.ServiceModel.Channels.AddressHeader[] { });

                    if (this.tracingEnabled)
                    {
                        Trace.Write(String.Format("Configured endpoints with SPN Identity claim.\nSPN Identity: {0}\nFIM MEX Endpoint: {1}\nFIM Enumeration Endpoint: {2}", fimServiceSPN, mexEndpointUri.AbsoluteUri, enumerationEndpointUri.AbsoluteUri));
                    }
                }

                // Initialize MEX client.
                this.mexClient = new MexClient(mexBinding, mexEndpoint);

                // Initialize WSEnumeration client.
                this.enumerationClient = new WsEnumerationClient(enumerationBinding, enumerationEndpoint);

                // If username, password and domain is specified, create and set client credential otherwise ADFS service account will be used.
                if (!string.IsNullOrEmpty(userName) && !string.IsNullOrEmpty(password) && !string.IsNullOrEmpty(userDomain))
                {
                    NetworkCredential credential = new NetworkCredential(userName, password, userDomain);
                    this.enumerationClient.ClientCredentials.Windows.ClientCredential = credential;
                    this.mexClient.ClientCredentials.Windows.ClientCredential = credential;

                    if (this.tracingEnabled)
                    {
                        Trace.Write(String.Format("Configured endpoints with the following credentials:\nDomain: {0}\nUserName: {1}", credential.Domain, credential.UserName));
                    }
                }
                else if (this.tracingEnabled)
                {
                    Trace.Write("Configured endpoints with ADFS Service credentials: " + WindowsIdentity.GetCurrent().Name);
                }
                
                // Get FIM schema and initialize Request and Resource factories.
                this.GetFIMSchema();
            }
            catch (Exception e)
            {
                AddLogException("Exception occured during initialization! ", e);
                throw;
            }

            if (this.tracingEnabled)
            {
                Trace.Write("XPath Attribute Store successfully initialized.");
            }
        }

        /// <summary>
        /// Called by ADFS to invoke asynchrounous query execution against Forefront Identity Manager.
        /// </summary>
        /// <param name="query">XPath query for selecting resources to retrieve from FIM</param>
        /// <param name="parameters">Attributes to retrieve from FIM resources.</param>
        /// <param name="callback">Reference to a method to be called when asynchronous operation completes.</param>
        /// <param name="state">The state information associated with the asynchronous operation.</param>
        /// <returns>The results of the asynchronous query against FIM.</returns>
        public IAsyncResult BeginExecuteQuery(string query, string[] parameters, AsyncCallback callback, object state)
        {
            try
            {
                string formattedQuery = string.Format(query, parameters);

                if (this.tracingEnabled)
                {
                    Trace.Write("XPath query to execute:\n" + formattedQuery);
                }

                AsyncResult queryResult = new TypedAsyncResult<string[][]>(callback, state);
                RunQueryDelegate queryDelegate = new RunQueryDelegate(this.RunQuery);
                queryDelegate.BeginInvoke(formattedQuery, new AsyncCallback(this.AsyncQueryCallback), queryResult);
                return queryResult;
            }
            catch (Exception e)
            {
                AddLogException("Exception occured in BeginExecuteQuery! ", e);
                throw;
            }
        }

        /// <summary>
        /// Method called by ADFS to get the query results.
        /// </summary>
        /// <param name="result">IAsyncResult reference that is returned by the BeginExecuteQuery method</param>
        /// <returns>The results of the asynchronous query against FIM</returns>
        /// <exception cref="Exception">Any Exception thrown during the asynchronous operation.</exception>
        public string[][] EndExecuteQuery(IAsyncResult result)
        {
            try
            {
                /* Each column of the returned array corresponds to a column requested in the query. Each row in the array corresponds to a row of 
                 * data within the requested columns. For example, for the query "Age=33;EmpName,Role", if there are two rows that match Age=33, 
                 * then the returned data will be:
                string[0][0] = "John", string[0][1] = "HR Manager"
                string[1][0] = "Leann",  string[1][1] = "Admin" */

                // Get results an throw any exception thrown during asynchronous operation.
                string[][] array = TypedAsyncResult<string[][]>.End(result);

                if (this.tracingEnabled)
                {
                    string str = string.Join(string.Empty, array.Select(subArray => "[" + string.Join("][", subArray) + "]\n").ToArray());
                    if (str.Length > 10000)
                    {
                        str = str.Substring(0, 10000) + "...";
                    }

                    Trace.Write("Issuing the following values:\n" + str);
                }
                
                return array;
            }
            catch (Exception e)
            {
                AddLogException("Exception occured in EndExecuteQuery! ", e);
                throw;
            }
        }

        #endregion

        #region IDisposable members

        /// <summary>
        /// Method to close or release unmanaged resources.
        /// </summary>
        public void Dispose()
        {
            this.Dispose(true);
            GC.SuppressFinalize(this);
        }

        #endregion
        
        /// <summary>
        /// Prepares and validates XPath query and attribute list for FIM.
        /// </summary>
        /// <param name="formattedQuery">Query and attributes concatenated from ADFS.</param>
        /// <param name="xpathQuery">XPath query</param>
        /// <param name="attributeNames">The attribues as array.</param>
        /// <exception cref="AttributeStoreQueryFormatException">Thrown when invalid query is coming from ADFS.</exception>
        private static void PrepareQuery(string formattedQuery, out string xpathQuery, out string[] attributeNames)
        {
            // We have two parts of the query - The XPath and the attributes to return separated by semicolon, 
            // for example: "/Person[AccountName='{0}'];ObjectID, DisplayName, EmployeeID".
            string[] queryParts = formattedQuery.Split(new char[] { ';' });
            if (queryParts.Length != 2)
            {
                throw new AttributeStoreQueryFormatException("Invalid query, the query must be made out of two parts - XPath query and attributes to return. " + formattedQuery);
            }

            // Trim the query parts.
            xpathQuery = queryParts[0].Trim();
            string attributes = queryParts[1].Trim();

            // Check if the query parts are valid.
            if (xpathQuery.Length == 0)
            {
                throw new AttributeStoreQueryFormatException("Empty XPath query. " + formattedQuery);
            }

            if (attributes.Length == 0)
            {
                throw new AttributeStoreQueryFormatException("Empty list of attributes to return. " + formattedQuery);
            }

            // Split and trim the attribute names that should be returned.
            List<string> attributesList =
                attributes
                    .Split(new char[] { ',' })
                    .ToList();

            attributeNames =
                attributesList
                    .Select(attributeName => attributeName.Trim())
                    .ToArray();            
        }

        /// <summary>
        /// Extracts query results from FIM and formats them into two-dimensional array for use with ADFS.
        /// </summary>
        /// <param name="attributeNames">Attribute names requested in query to FIM.</param>
        /// <param name="resources">Resources returned from FIM.</param>
        /// <returns>FIM Query Results as two-dimensional string array where the columns in the array represent claim types, and the rows represent claim values.</returns>
        private static string[][] ExtractQueryResults(List<string> attributeNames, List<RmResource> resources)
        {
            // Data should be returned in the following format: data[rows][columns].
            // Create List(rows) of Lists(columns) to hold our data.
            List<List<string>> rows = new List<List<string>>();

            // Iterate resources.
            foreach (RmResource resource in resources)
            {
                // Create a "main" row for each resource.
                AddResultRow(attributeNames.Count, rows);

                // Iterate columns.
                for (int columnIndex = 0; columnIndex < attributeNames.Count; columnIndex++)
                {
                    RmAttributeValue attributeValue = resource.Attributes[new RmAttributeName(attributeNames[columnIndex])];
                    if (attributeValue != null)
                    {
                        if (!attributeValue.IsMultiValue)
                        {
                            // Single-valued attribute.
                            RmAttributeValueSingle singleValue = (RmAttributeValueSingle)attributeValue;
                            if (singleValue.Value != null)
                            {
                                rows.Last()[columnIndex] = attributeValue.ToString();
                            }
                        }
                        else
                        {
                            // Multi-valued attribute.
                            RmAttributeValueMulti multiValue = attributeValue as RmAttributeValueMulti;

                            // Iterate value within multi-value.
                            foreach (IComparable value in attributeValue.Values)
                            {
                                if (value != null)
                                {
                                    if (value != attributeValue.Values.First())
                                    {
                                        AddResultRow(attributeNames.Count, rows);
                                    }

                                    rows.Last()[columnIndex] = value.ToString();
                                }
                            }
                        }
                    }
                }
            }

            // Convert a list of lists to an array of arrays. Clear as mud!
            string[][] queryResults =
                rows
                    .Select(detail => detail.ToArray())
                    .ToArray();

            return queryResults;
        }

        /// <summary>
        /// Initializes a new row (List) for use when extracting query results from FIM.
        /// </summary>
        /// <param name="columnLenght">The number of items that should be initialized in List</param>
        /// <param name="rows">Parent List of Lists to hold rows.</param>
        private static void AddResultRow(int columnLenght, List<List<string>> rows)
        {
            List<string> row = new List<string>(columnLenght);

            // Initialize row to contain null for each entry.
            for (int i = 0; i < columnLenght; i++)
            {
                row.Insert(i, null);
            }

            rows.Add(row);
        }

        #region Logging

        /// <summary>
        /// Appends information log entry to event log.
        /// </summary>
        /// <param name="msg">Message to log.</param>
        private static void AddLogInfo(string msg)
        {
            AddLog(EventLogEntryType.Information, msg);
        }

        /// <summary>
        /// Appends exception log entry to event log.
        /// </summary>
        /// <param name="msg">Message to log.</param>
        /// <param name="e">Exception to log.</param>
        private static void AddLogException(string msg, Exception e)
        {
            AddLog(EventLogEntryType.Error, string.Format("{0}\n{1}", msg, e.ToString()));
        }

        /// <summary>
        /// Writes to event log.
        /// </summary>
        /// <param name="eventType">Type of log entry to write.</param>
        /// <param name="msg">Message to write.</param>
        private static void AddLog(EventLogEntryType eventType, string msg)
        {
            try
            {
                EventLog.WriteEntry(LogSource, msg, eventType);
            }
            catch
            {
            }
        }

        #endregion

        /// <summary>
        ///  Callback on completed asynchronous call.
        /// </summary>
        /// <param name="result">The results of the asynchronous operation.</param>
        private void AsyncQueryCallback(IAsyncResult result)
        {
            TypedAsyncResult<string[][]> queryResult = (TypedAsyncResult<string[][]>)result.AsyncState;
            RemoteMessagingAsyncResult delegateAsyncResult = (RemoteMessagingAsyncResult)result;
            RunQueryDelegate runQueryDelegate = (RunQueryDelegate)delegateAsyncResult.AsyncDelegate;

            string[][] resultingValues = null;
            Exception originalException = null;
            try
            {
                resultingValues = runQueryDelegate.EndInvoke(result);
            }
            catch (Exception e)
            {
                // We don't want exceptions to be thrown from the callback method as these need to be made available to the thread that calls EndExecuteQuery.
                AddLogException("Exception occured in AsyncQueryCallback! ", e);
                originalException = e;
            }

            // Any exception is stored in query Result and re-thrown when EndExecuteQueryMethod calls TypedAsyncResult<string[][]>.End(..) method.
            queryResult.Complete(resultingValues, false, originalException);
        }

        /// <summary>
        ///  Executes query against FIM's WS Enumeration endpoint.
        /// </summary>
        /// <param name="formattedQuery">FIM XPath query and attributes from ADFS.</param>
        /// <returns>FIM Query Results as two-dimensional string array where the columns in the array represent claim types, and the rows represent claim values.</returns>
        /// <exception cref="AttributeStoreQueryFormatException">Thrown when invalid query is coming from ADFS.</exception>
        /// <exception cref="AttributeStoreQueryExecutionException">Thrown when query couldn't be executed against FIM.</exception>
        private string[][] RunQuery(string formattedQuery)
        {
            try
            {
                // Prepare query for FIM.
                string xpathQuery;
                string[] attributeNames;
                PrepareQuery(formattedQuery, out xpathQuery, out attributeNames);

                // Execute query against FIM.
                List<RmResource> resources =
                    this.requestFactory
                    .CreateEnumeration(this.enumerationClient, this.resourceFactory, xpathQuery, attributeNames)
                    .ToList();

                // Extracts query results into an array of arrays and return it.
                return ExtractQueryResults(attributeNames.ToList(), resources);
            }
            catch (AttributeStoreQueryFormatException e)
            {
                AddLogException("Exception occured in RunQuery! ", e);
                throw;
            }
            catch (Exception e)
            {
                // Wrap any unknown exceptions.
                AddLogException("Exception occured in RunQuery! ", e);
                throw new AttributeStoreQueryExecutionException("Query Execution Failed. ", e);
            }
        }

        /// <summary>
        /// Gets schema metadata from FIM and initailizes request and resource factories.
        /// </summary>
        /// <returns>FIM Schema as XmlSchemaSet.</returns>
        /// <exception cref="Exception">Exception if schema couldn't be recieved.</exception>
        private XmlSchemaSet GetFIMSchema()
        {
            XmlSchemaSet metadata = null;
            try
            {
                // Get schema metadata.
                metadata = this.mexClient.Get();

                // Initialize request and resource factories with metadata.
                this.requestFactory = new RmRequestFactory(metadata);
                this.resourceFactory = new RmResourceFactory(metadata);

                if (this.tracingEnabled)
                {
                    Trace.Write("Successfully collected FIM schema metadata.");
                }
            }
            catch (Exception e)
            {
                AddLogException("Exception occured when trying to get schema metadata from FIM! ", e);
                throw;
            }

            return metadata;
        }

        /// <summary>
        /// Method to close or release unmanaged resources.
        /// </summary>
        /// <param name="disposing">false if the method has been called by the runtime otherwise true.</param>
        private void Dispose(bool disposing)
        {
            if (!this.disposed)
            {
                if (disposing)
                {
                    this.traceListener.Dispose();
                    this.mexClient.Close();
                    this.enumerationClient.Close();
                }

                // Note disposing has been done.
                this.disposed = true;
            }
        }
    }
}
